from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy import Column, Integer, String, Float, DateTime
import datetime
import aiosqlite
from sqlalchemy import select

DATABASE_URL = "sqlite+aiosqlite:///./bookkeeping.db"

engine = create_async_engine(DATABASE_URL, echo=True, future=True)
SessionLocal = sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
Base = declarative_base()


class Bill(Base):
    __tablename__ = "bills"
    id = Column(Integer, primary_key=True, index=True)
    description = Column(String, index=True)
    amount = Column(Float)
    category = Column(String)
    date = Column(DateTime, default=datetime.datetime.utcnow)
    type = Column(String, default="支出")


class BillCreate(BaseModel):
    description: str
    amount: float
    category: str
    date: Optional[datetime.datetime] = None
    type: str = "支出"


class BillOut(BillCreate):
    id: int
    date: datetime.datetime
    type: str = "支出"

    class Config:
        from_attributes = True


async def lifespan(app: FastAPI):
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield

app = FastAPI(lifespan=lifespan)

# 允许前端跨域
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.post("/bills/", response_model=BillOut)
async def create_bill(bill: BillCreate):
    async with SessionLocal() as session:
        db_bill = Bill(**bill.dict())
        session.add(db_bill)
        await session.commit()
        await session.refresh(db_bill)
        return db_bill


@app.get("/bills/", response_model=List[BillOut])
async def read_bills(skip: int = 0, limit: int = 100):
    async with SessionLocal() as session:
        stmt = select(Bill).offset(skip).limit(limit)
        result = await session.execute(stmt)
        bills = result.scalars().all()
        return bills


@app.get("/bills/{bill_id}", response_model=BillOut)
async def read_bill(bill_id: int):
    async with SessionLocal() as session:
        result = await session.get(Bill, bill_id)
        if not result:
            raise HTTPException(status_code=404, detail="账目未找到")
        return result


@app.put("/bills/{bill_id}", response_model=BillOut)
async def update_bill(bill_id: int, bill: BillCreate):
    async with SessionLocal() as session:
        db_bill = await session.get(Bill, bill_id)
        if not db_bill:
            raise HTTPException(status_code=404, detail="账目未找到")
        for key, value in bill.dict().items():
            setattr(db_bill, key, value)
        await session.commit()
        await session.refresh(db_bill)
        return db_bill


@app.delete("/bills/{bill_id}")
async def delete_bill(bill_id: int):
    async with SessionLocal() as session:
        db_bill = await session.get(Bill, bill_id)
        if not db_bill:
            raise HTTPException(status_code=404, detail="账目未找到")
        await session.delete(db_bill)
        await session.commit()
        return {"ok": True}


@app.get("/export/")
async def export_bills():
    async with SessionLocal() as session:
        result = await session.execute(Base.metadata.tables["bills"].select())
        bills = result.fetchall()
        # 导出为CSV格式
        import csv
        from fastapi.responses import StreamingResponse
        from io import StringIO

        output = StringIO()
        writer = csv.writer(output)
        writer.writerow(["id", "description", "amount", "category", "date"])
        for row in bills:
            writer.writerow(
                [row.id, row.description, row.amount, row.category, row.date]
            )
        output.seek(0)
        return StreamingResponse(
            output,
            media_type="text/csv",
            headers={"Content-Disposition": "attachment; filename=bills.csv"},
        )
